from langchain_ollama.embeddings import OllamaEmbeddings
from vector_database import chunk
import os

def embed(text: str) -> list[float]:
    """
        take a str as para, and return the vector ([float]) only for this str
    """
    # get the embedded vectors
    model = os.getenv("OLLAMA_EMBEDDING_MODEL") or "nomic-embed-text"
    embeddings_model = OllamaEmbeddings(model=model)
    print(f"→ Converting \"{text[:60].replace("\n", " ")}...\" ", end="")
    embeddings = embeddings_model.embed_documents(
        [text]
    )
    print("✔ Done")

    return embeddings[0]


def embed_bulk(texts: [str]) -> list[list[float]]:
    total = len(texts)
    result = []
    for i, text in enumerate(texts):
        print(f"({i + 1}/{total}) ", end="")
        result.append(
            embed(text)
        )
    return result


if __name__ == "__main__":
    source_text = chunk.get_chunks_content_only()[:5]
    vectors = embed_bulk(source_text)
    for v in vectors:
        print("=" * 40)
        print(v)
